package com.isaiah.tree;

import com.isaiah.list.ArrayList;

/**
 * @Title: 自平衡二叉查找树
 * @Description: ${todo}
 * @author: Haijun
 * @emial: isaiah@sohu.com
 * @date 2020/4/817:01
 */
public class AVLTree<K extends Comparable<K>, V> {
    private class Node {
        private K key;
        private V value;
        public Node left, right;
        public int height;

        public Node(K key, V value) {
            this.key = key;
            this.value = value;
            this.left = null;
            this.right = null;
            height = 1;     // 默认新的一个节点是叶子节点，高度为1

        }
    }

    private Node root;
    private int size;

    public AVLTree(){
        this.root = null;
        this.size = 0;
    }

    public int getSize(){
        return size;
    }

    public boolean isEmpty(){
        return size == 0;
    }

    /**
     * 获取指定节点的高度
     * @param node
     * @return
     */
    private int getHeight(Node node){
        if (node == null){
            return 0;
        }
        return node.height;
    }

    /**
     * 计算平衡因子
     * @param node
     * @return
     */
    private int getBalanceFactor(Node node){
        if (node == null){
            return 0;
        }
        // 左子树的高度-右子树的高度
        return getHeight(node.left) - getHeight(node.right);
    }

    /**
     * 判断该二叉树是不是一棵二分搜索树
     * @return
     */
    public boolean isBST(){
        ArrayList<K> keys = new ArrayList<>();
        // 使用中序遍历进行遍历整个树，然后判断遍历后的结果是不是从小到大的顺序排列
        inOrder(root, keys);
        for (int i = 0; i < keys.getSize() -1; i++){
            if (keys.get(i).compareTo(keys.get(i+1)) > 0){
                return false;
            }

        }
        return true;
    }

    /**
     * 中序遍历
     * @param node
     * @param keys
     */
    private void inOrder(Node node, ArrayList<K> keys) {
        if (node == null){
            return ;
        }

        inOrder(node.left, keys);
        keys.addLast(node.key);
        inOrder(node.right, keys);
    }

    /**
     * 判断该二叉树是不是平衡二叉树
     * @return
     */
    public boolean isBalanced(){
        return isBalanced(root);
    }

    private boolean isBalanced(Node node){
        if (node == null){
            return true;
        }
        // 得到节点的平衡因子
        int balanceFactor = getBalanceFactor(node);
        if (Math.abs(balanceFactor) > 1){
            return false;
        }
        // 判断左子树与右子树是否平衡
        return isBalanced(node.left) && isBalanced(node.right);

    }


    /**
     * 添加节点
     * @param key
     * @param value
     */
    public void add(K key, V value){
        root = add(root, key, value);
    }

    private Node add(Node node, K key, V value){
        if (node == null){
            size ++;
            return new Node(key, value);
        }

        if (key.compareTo(node.key)< 0){        // 添加到左边
            node.left = add(node.left, key, value);
        }else if (key.compareTo(node.key) > 0){
            node.right = add(node.right, key, value); // 添加右边
        }else{  //key.compareTo(node.key) == 0

            node.value = value;                 // 覆盖
        }
        // 维护高度
        node.height = 1 + Math.max(getHeight(node.left), getHeight(node.right));

        // 计算平衡因子
        int balanceFactor = getBalanceFactor(node);


        // 情况一：LL 左子树高度比右子树高度超过1, 左子树平衡因子大于等于0 （当前节点左侧多添加了一个节点）
        if (balanceFactor > 1 && getBalanceFactor(node.left) >= 0){
            return rightRotate(node);
        }

        // 情况二：RR  右子树高度比左子树高差超过1, 右子树平衡因子大于等于0 （当前节点右侧多添加了一个节点）
        if (balanceFactor < -1 && getBalanceFactor(node.right) <= 0){
            return leftRotate(node);
        }

        // 情况三：LR  左子树比右子树高且高度超过1，对于Node的左子树的它的右子树比左子树还要高
        if (balanceFactor >1 && getBalanceFactor(node.left) < 0){
            // 1)对当前节点左孩子进行左旋转
            node.left = leftRotate(node.left);
            // 2) 对当前节点进行右旋转
            return rightRotate(node);

        }
        // 情况四：LR  右子树比左子树高且高度差超过1，对于Node的右子树的它的左子树比右子树还要高
        if (balanceFactor < -1 && getBalanceFactor(node.right) > 0){
            // 1)对当前节点右孩子进行右旋转
            node.right = rightRotate(node.right);
            // 2) 对当前节点进行左旋转
            return leftRotate(node);

        }

        return node;
    }

    public V get(K key){
        // 1. 根据key查找节点
        Node node = getNode(root, key);
        // 2. 返回node的value
        return node == null? null : node.value;
    }


    /**
     * 删除任意元素
     * @param key
     * @return
     */
    public V remove(K key){

        Node node = getNode(root, key);
        if(node != null){
            root = remove(root, key);
            return node.value;
        }
        return null;
    }


    /**
     * 从指定的node为根的子树中删除元素
     * @param node
     * @param key
     * @return
     */
    private Node remove(Node node,  K key) {
        if (node == null){
            return null;
        }
        Node retNode;
        if (key.compareTo(node.key) < 0){       // 说明要删除的元素在左子树中
            node.left = remove(node.left, key);
            retNode =  node;
        }
        else if (key.compareTo(node.key) > 0){  // 说明要删除的元素在右子树中
            node.right = remove(node.right, key);
            retNode =  node;
        }
        else{// e == node.e                 // 说明要删除的元素正是当前元素

            // 待删除节点左子树为null的情况
            if (node.left == null){
                Node rightNode = node.right;
                node.right = null;
                size --;
                retNode =  rightNode;
            }
            // 待删除节点右子树为null的情况
            else if (node.right == null){
                Node leftNode = node.left;
                node.left = null;
                size --;
                retNode =  leftNode;
            }
            else{
                /** Hibbard 删除逻辑 */
                // 1. 找到比待删除元素大的最小节点：即待删除节点右子树中最小节点
                Node successor = mininum(node.right);
                // 2. 让待删除节点右子树中最小节点right指向删除节点右子树中最小节点后的子树
                //successor.right = removeMin(node.right);
                successor.right = remove(node.right, successor.key);
                // 3. 让待删除节点右子树中最小节点left指向原来left
                successor.left = node.left;
                // 4. 删除节点
                node.left = node.right = null;
                retNode =  successor;
            }

        }
        if (retNode == null){
            return null;
        }
        // 更新height
        retNode.height =  Math.max(getHeight(retNode.left), getHeight(retNode.right)) + 1;

        // 计算平衡因子
        int balanceFactor = getBalanceFactor(retNode);

        // 平衡维护
        // LL
        if (balanceFactor > 1 && getBalanceFactor(retNode.left) >= 0){
            return rightRotate(retNode);
        }
        // RR
        if (balanceFactor <-1 && getBalanceFactor(retNode.right) <= 0){
            return leftRotate(retNode);
        }

        // LR
        if (balanceFactor >1 && getBalanceFactor(retNode.left) < 0){
            retNode.left = leftRotate(retNode.left);
            return rightRotate(retNode);
        }
        // RL
        if (balanceFactor < -1 && getBalanceFactor(retNode.right) > 0){
            retNode.right = rightRotate(retNode.right);
            return leftRotate(retNode);
        }
        return retNode;
    }
    /**
     * 查找二分搜索树的最小元素
     * @return
     */
    public Node mininum(){
        if (size == 0){
            throw new IllegalArgumentException("BST is empty！");
        }
        return  mininum(root);
    }
    /**
     * 返回以node为根的二分搜索树的最小值所在的节点
     * @param node
     * @return
     */
    public Node mininum(Node node){
        if (node.left == null){
            return node;
        }
        // 一直搜索左子树
        return mininum(node.left);
    }

    /**
     * 删除二分搜索树的最小元素
     * @return
     */
    public Node removeMin(){
        // 找到最小元素
        Node ret = mininum();
        // 删除最小元素
        root = removeMin(root);
        // 将最小元素返回
        return ret;
    }


    // 看二分搜索树中是否包含元素e
    public boolean contains(K key){
        return contains(root, key);
    }

    // 看以node为根的二分搜索树中是否包含元素e, 递归算法
    private boolean contains(Node node, K key){

        if(node == null){
            return false;
        }

        if(key.compareTo(node.key) == 0){
            return true;
        }
        else if(key.compareTo(node.key) < 0){
            return contains(node.left, key);
        }
        else { // e.compareTo(node.e) > 0
            return contains(node.right, key);
        }

    }

    /**
     * 更新key对应的value
     * @param key
     * @param newValue
     */
    public void set(K key, V newValue){
        // 1. 查找到节点
        Node node = getNode(root, key);
        if (node == null){
            throw new IllegalArgumentException(key + " doesn't exist!");
        }
        // 2. 更新节点value
        node.value = newValue;
    }

    private Node getNode(Node node, K key){
        if (node == null){
            return null;
        }
        if (key.equals(node.key)){
            return node;
        }
        // 在左子树中查找
        else if (key.compareTo(node.key) < 0){
            return getNode(node.left, key);
        }
        // 在右子树中查找
        else { // key.compareTo(node.key) >0
            return getNode(node.right, key);
        }
    }

    private Node removeMin(Node node) {
        if (node.left == null){
            // 当前节点可能有右孩子
            Node rightNode = node.right;
            // 把要删除的node的右孩子置为null
            node.right = null;
            size--;
            return rightNode;
        }
        node.left = removeMin(node.left);
        return node;
    }
    /**
     * 对节点y进行向右旋转操作，返回旋转后新的根节点x
     *              y                                   x
     *             / |                                /   \
     *            x   T4         向右旋转(y)          z     y
     *           / \         ------------------>    / \   / \
     *          z   T3                             T1 T2 T3 T4
     *         / |
     *        T1  T2
     * @param y
     * @return
     */
    private Node rightRotate(Node y){
        Node x = y.left;
        Node T3 = x.right;
        // 向右旋转过程
        x.right = y;
        y.left = T3;
        // 更新高度
        y.height = Math.max(getHeight(y.left), getHeight(y.right)) + 1;
        x.height = Math.max(getHeight(x.left), getHeight(x.right)) + 1;
        return x;
    }

    /**
     * 对节点y进行向左旋转操作，返回旋转后新的根节点x
     *               y                                   x
     *              / |                                /   \
     *             T1  x         向左旋转(y)            y    z
     *                / |        ------------------>  / \   / \
     *              T2  z                           T1 T2 T3 T4
     *                 / |
     *                T3 T4
     * @param y
     * @return
     */
    private Node leftRotate(Node y){
        Node x = y.right;
        Node T2 = x.left;
        // 向左旋转过程
        x.left = y;
        y.right = T2;
        // 更新高度
        y.height = Math.max(getHeight(y.left), getHeight(y.right)) + 1;
        x.height = Math.max(getHeight(x.left), getHeight(x.right)) + 1;
        return x;
    }


}
